[mlir][rocdl] Add readfirstlane intrinsic
#152551
Conversation
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-llvm Author: Ivan Butygin (Hardcode84) ChangesFull diff: https://github.com/llvm/llvm-project/pull/152551.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index a2354e22e2745..992dc2209372c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -189,7 +189,21 @@ def ROCDL_BallotOp :
let assemblyFormat = "$pred attr-dict `:` type($res)";
}
-def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>], 1>,
+def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0], [AllTypesMatch<["res", "src"]>, Pure], 1>,
+ Arguments<(ins LLVM_Type:$src)> {
+ let results = (outs LLVM_Type:$res);
+ let summary = "Get the value in first active lane.";
+
+ let description = [{
+ Returns the value in the lowest active lane of the input operand.
+ }];
+
+ let assemblyFormat = [{
+ $src attr-dict `:` type($res)
+ }];
+}
+
+def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>, Pure], 1>,
Arguments<(ins LLVM_Type:$src0,
I32:$src1)> {
let results = (outs LLVM_Type:$res);
@@ -201,7 +215,7 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res",
let assemblyFormat = [{
$src0 `,` $src1 attr-dict `:` `(` type($src0) `,` type($src1) `)` `->` type($res)
- }];
+ }];
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index a2b2f84606ba0..db5271c57f573 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -981,6 +981,13 @@ llvm.func @rocdl.s.wait.expcnt() {
// -----
+llvm.func @rocdl.readfirstlane(%src : f32) -> f32 {
+ // CHECK-LABEL: rocdl.readfirstlane
+ // CHECK: rocdl.readfirstlane %{{.*}} : f32
+ %ret = rocdl.readfirstlane %src : f32
+ llvm.return %ret : f32
+}
+
llvm.func @rocdl.readlane(%src : f32) -> f32 {
%cst0 = llvm.mlir.constant(0 : i32) : i32
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 740990a6e589b..ce4394102b5a1 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -125,6 +125,23 @@ llvm.func @rocdl.ballot64(%pred : i1) -> i64 {
llvm.return %0 : i64
}
+llvm.func @rocdl.readfirstlane(%src0 : f32, %src1: f64, %src2: i32, %src3: vector<2 x f32>) -> f32 {
+ // CHECK-LABEL: rocdl.readfirstlane
+ // CHECK: call float @llvm.amdgcn.readfirstlane.f32(float %{{.*}})
+ %0 = rocdl.readfirstlane %src0 : f32
+
+ // CHECK: call double @llvm.amdgcn.readfirstlane.f64(double %{{.*}})
+ %1 = rocdl.readfirstlane %src1 : f64
+
+ // CHECK: call i32 @llvm.amdgcn.readfirstlane.i32(i32 %{{.*}})
+ %2 = rocdl.readfirstlane %src2 : i32
+
+ // CHECK: call <2 x float> @llvm.amdgcn.readfirstlane.v2f32(<2 x float> %{{.*}})
+ %3 = rocdl.readfirstlane %src3 : vector<2 x f32>
+
+ llvm.return %0 : f32
+}
+
llvm.func @rocdl.readlane(%src0 : f32, %src1: f64, %src2: i32, %src3: vector<2 x f32>) -> f32 {
%idx = llvm.mlir.constant(0 : i32) : i32
|
| @@ -189,7 +189,21 @@ def ROCDL_BallotOp : | |||
| let assemblyFormat = "$pred attr-dict `:` type($res)"; | |||
| } | |||
|
|
|||
| def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>], 1>, | |||
| def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0], [AllTypesMatch<["res", "src"]>, Pure], 1>, | |||
There was a problem hiding this comment.
IIUC, this relies on the execution mask to determine the first active lane. Do we really want this to be pure and enable speculation/hoisting, since hoisting outside of ifs/loops can affect active lanes?
There was a problem hiding this comment.
Ok, making it loop-hoistable was my exact motivation, any ideas how to handle it?
There was a problem hiding this comment.
I guess, I can make it non-pure for now and hoist manually on Wave side.
There was a problem hiding this comment.
I think you can only hoist it if you can prove the control flow is uniform
There was a problem hiding this comment.
Yeah, I see. Can we hoist readlane?
There was a problem hiding this comment.
I think structured control flow guarantees that your parents are at least as active as you, so this seems safe for me
There was a problem hiding this comment.
Yeah, this isn't hoistable in general, but you can slide out of a for loop or some other thing like that. That is, this isn't speculatable, but it does have no memory effects
|
I will drop pure completely for now, as I can do hoisting on our side completely. |
No description provided.